#!/usr/bin/python

#RB deleted unused functions
#RB neutralize molecues via neutralisation function in chemistry module
#RB changed function for non-defined linkers
#RB if same CF occurs several times in same moli => only store once

import string,os, sys, mysql
#import os,sys,string,re
from openeye.oechem import *
import chemistry

username = 'webuser'
password = ''

smiles_file =  oemolistream(sys.argv[1])
#smiles_file=open(sys.argv[2], 'r')
out_file=open(sys.argv[2],'w') 
out_file_failed = open(sys.argv[3],'w')
kind_of_fragments = sys.argv[4]

conn=mysql.connect2server(password, username, 'purchasable')  

debug = False


Hal = ['I', 'Br', 'Cl', 'F', 'H']
Other_atoms = ['C', 'S', 'N', 'P']
#--------------------------------------------
def get_rules(password, username, kind_of_rules):			  
	cursor = conn.cursor ()
	if kind_of_rules == 'standard':
           command = "select smirks,description,name from purchasable.definition_make_scaffolds where name like '%standard%' order by id"
        elif kind_of_rules == 'after':
           command = "select smirks,description,name from purchasable.definition_make_scaffolds where name like '%after%' order by id"
        elif kind_of_rules == 'simple':
           command = "select smirks,description,name from purchasable.definition_make_scaffolds wnhere name like '%simple%' order by id"

	rule = []
	
        
        #if kind_of_fragments == 'ring':
	   #command = 'select smirks,name from definition_make_scaffolds where id = 191 or id = 192 or id =193 or id =194 or id =195 or id =196 or id =197  order by id'
	print command
	cursor.execute(command)	
	rules = cursor.fetchall()
        cursor.close ()
	k = 0
	for i in rules:
		rule.append([i[0],i[1]])
		#print rule
		k=k+1
	return rule


#--------------------------------------------

def split_fragments(split_mol,mol_list):
	iso = False #keep stereoinformation
	kek = False #kekulize to get unique smiles (no, this is no longer true)
	chiral = False #keep information about chirality (could also be false, because this information is guessed 	
	non_ring_fragments = False
	

	#the two next lines are important to make things more consistent with daylight kit
	OEFindRingAtomsAndBonds(split_mol)
	OEAssignAromaticFlags(split_mol)
	if OECreateCanSmiString(split_mol).find('.') <> -1:  #returns -1 if '.' is not found, '<>' = '!='
		#remove part wo ring
		partcount,partlist=OEDetermineComponents(split_mol)
		#print partlist
		#print partcount,partlist
		for i in xrange(1,partcount+1):
			#print 'check part ' + str(i)
			pred = OEPartPredAtom(partlist) # ?
			test_mol = OECreateOEMol()      # ?
			pred.SelectPart(i)             
			OESubsetMol(test_mol, split_mol, pred) # ?
			#print 'test ' + OECreateCanSmiString(test_mol)
			rings, ringlist = OEDetermineRingSystems(test_mol)
			if rings >= 0:
				mol_list.append(test_mol)
				#print 'found part with ring ' + OECreateCanSmiString(test_mol) 
			else:
				non_ring_fragments = True
	else:
		mol_list.append(split_mol)

	clean_list = []
	for i in mol_list:
	   	#this strange pattern is necessary to make a unique smiles, this is only possible after the clean-up
		OEFindRingAtomsAndBonds(i)
     	 	OEAssignAromaticFlags(i)
	 	smi = chemistry.CanSmi(i,iso,kek,chiral)
         	smi = chemistry.clean_smiles(smi)
		#print smi + '<------------------------'
	 	clean_mol = OEGraphMol()
	 	OEParseSmiles(clean_mol, smi)	
		clean_list.append(clean_mol)
	mol_list = clean_list
	return mol_list


	
#--------------------------------------------
def transform_framework(rules,mainframework,title,debug,more_rings):
	iso = False #keep stereoinformation
	kek = False #kekulize to get unique smiles (no, this is no longer true)
	chiral = False #keep information about chirality (could also be false, because this information is guessed anyway)
	if debug:
		print '-----------> start'
	failed = False
	more_rings = more_rings
	scaffold_list = []
	transform_mol = OEGraphMol()
        #if test == False:
	OEParseSmiles(transform_mol, mainframework)
	#if test == True:
      
	OEFindRingAtomsAndBonds(transform_mol)
	OEAssignAromaticFlags(transform_mol)
	OEAssignFormalCharges(transform_mol)

	scaffold_list.append(transform_mol)
  	#old_smi = OECreateCanSmiString(transform_mol)
	if more_rings == True:
		#print rules[1], ' rules[0]'
		ij = 0
		counter = False
		while ij < len(rules) and counter != True: 
			rule = rules[ij]
		#while len(scaffold_list)  <2:
  		#for rule in rules:
			#while len(scaffold_list)  <2:
	 	  	mol_list = []
			for transform_mol in scaffold_list:
				old_smi = OECreateCanSmiString(transform_mol)
	  	  		umr = OEUniMolecularRxn(rule[0])
				umr(transform_mol)
				if string.count(OECreateCanSmiString(transform_mol), '.') == 1:
					counter = True
				if string.count(OECreateCanSmiString(transform_mol), '.') >= 2:
					counter = True
					failed = True
					
				old_len_mol_list = len(mol_list)
				if debug:
					if old_smi <> OECreateCanSmiString(transform_mol):
						print '--->'
						print rule[0], '--', rule[1]
						print old_smi
						old_smi = OECreateCanSmiString(transform_mol)
						print old_smi
		
				mol_list = split_fragments(transform_mol,mol_list)

	 		scaffold_list = []

	  	 	for i in mol_list:
				count_ring, ringlist = OEDetermineRingSystems(i)
				if count_ring >= 1:
					scaffold_list.append(i)
			ij = ij +1
	else:
		for rule in rules:
	   		mol_list = []
       	  		for transform_mol in scaffold_list:
	
				old_smi = OECreateCanSmiString(transform_mol)
	  	  		umr = OEUniMolecularRxn(rule[0])
				umr(transform_mol)
				old_len_mol_list = len(mol_list)
		
				if debug:
					if old_smi <> OECreateCanSmiString(transform_mol):
						print '--->', rule[0], '<---', rule[1]
						print old_smi
						old_smi = OECreateCanSmiString(transform_mol)
						print old_smi
						print
		
				mol_list = split_fragments(transform_mol,mol_list)

	 		scaffold_list = []

	  	 	for i in mol_list:
				count_ring, ringlist = OEDetermineRingSystems(i)
				if count_ring >= 1:
					scaffold_list.append(i)
			
	if debug:
		print 'generated the following scaffolds:'
		for i in scaffold_list:
			print OECreateCanSmiString(i)
	if len(scaffold_list) == 1 and more_rings == True:
		print 'hier rein'
		failed = True
	if failed:
		scaffold_list[0] ='failed'
	return scaffold_list
#-----------------------
def break_linkers(mol,count_ring,ring_list):

	smiles_list = []
	number = count_ring 
	anchor_dict = {}  # part_number: [ring number, anchor atom index]
	for i in range(0,len(ring_list)):
			if  ring_list[i] == 0 or ring_list[i] > count_ring:
				#check atom
				number = number + 1
				swap_number_dict = {}
				atom = mol.GetAtom(OEHasAtomIdx(i))
				if debug: print OEGetAtomicSymbol(atom.GetAtomicNum()), '<-----'
				ring_list[i] = number
				#check the neighbours
				if debug:
					print 'check neighbours'
				for neigh in atom.GetAtoms():
					if debug:
						print 'neighbour: ', OEGetAtomicSymbol(neigh.GetAtomicNum())
						print ring_list[neigh.GetIdx()]
					if ring_list[neigh.GetIdx()] == 0: #neighbour not yet assigend to anything => same catogery than orginial atom
						ring_list[neigh.GetIdx()] = number
					elif ring_list[neigh.GetIdx()] > count_ring: #already part of linker => mark number of linker
						swap_number_dict[ring_list[neigh.GetIdx()]] = number
					elif anchor_dict.has_key(number): #neighbour is part of ring => join ring number and ring atom with linker number
						anchor_dict[number].append([ring_list[neigh.GetIdx()],neigh.GetIdx()]) 
					else:
						anchor_dict[number]= [[ring_list[neigh.GetIdx()],neigh.GetIdx()]]


					if debug: print anchor_dict
			     	if len(swap_number_dict) > 0:
					#swap numbers
					for j in range(0,len(ring_list)):
						if swap_number_dict.has_key(ring_list[j]):
							if anchor_dict.has_key(ring_list[j]):
								if anchor_dict.has_key(swap_number_dict[ring_list[j]]):
									for index in anchor_dict[ring_list[j]]:
									 anchor_dict[swap_number_dict[ring_list[j]]].append(index)
								else:
									anchor_dict[swap_number_dict[ring_list[j]]] = anchor_dict[ring_list[j]]
							anchor_dict[ring_list[j]] = []
							ring_list[j] =swap_number_dict[ring_list[j]]
				if debug: print ring_list
							

	split_dict = {}
	for key in anchor_dict:
		if len(anchor_dict[key]) == 1:
			#terminal group
			for j in range(0,len(ring_list)):
				if ring_list[j] == key:
					ring_list[j] = anchor_dict[key][0][0]
			if debug: print ring_list, 'joined terminal groups' 
		elif len(anchor_dict[key]) > 1:
			#linker
			split_dict[key] = anchor_dict[key]
	if debug: print split_dict, 'split dict'

	#prepeare dict with list for each ring (already joined with terminal groups) and moli which can be modified
	result_dict = {}
	for ring in range(1,count_ring + 1):
		temp_list = []
		for ring_atom in range(0, len(ring_list)):
			if ring_list[ring_atom] ==ring:
				temp_list.append(ring)
			else:
				temp_list.append(0)

		new_mol = OECreateOEMol()
		new_mol = mol
		result_dict[ring] = [temp_list,new_mol]
	#if debug: print result_dict, 'result_dict'	
	

	for linker in split_dict:
		for ring_num, ring_atom in split_dict[linker]:

			for j in range(0, len(ring_list)):
				if ring_list[j] == linker: #add linker to current ring
					result_dict[ring_num][0][j]=ring_num

			#add ring_atom that is attached to linker to new core fragment
			for linked_ring, linked_ring_num in split_dict[linker]:
				if result_dict[ring_num][0][linked_ring_num] <> ring_num: # add linked atom to current ring
					atom = result_dict[ring_num][1].GetAtom(OEHasAtomIdx(linked_ring_num))
					atom.SetAromatic(False)
					result_dict[ring_num][0][linked_ring_num] = ring_num

			if debug: print result_dict[ring_num][0], 'modified ring_list', linker, 'linker'
	for ring_num in result_dict:

			pred = OEPartPredAtom(result_dict[ring_num][0]) 
			pred.SelectPart(ring_num)
			
			temp_mol = OECreateOEMol()
			OESubsetMol(temp_mol, result_dict[ring_num][1], pred)

     			OEFindRingAtomsAndBonds(temp_mol)
     			OEAssignAromaticFlags(temp_mol)
     			OEAssignFormalCharges(temp_mol)
     			new_smi = chemistry.CanSmi(temp_mol,False,False,False)
			new_smi = chemistry.clean_smiles(new_smi)
			if debug:
				print new_smi
			smiles_list.append(new_smi)

	return smiles_list	
		
				
						
	
#-----------------------
		


#Main


iso = False #keep stereoinformation
kek = False #kekulize to get unique smiles (no, this is no longer true)
chiral = False #keep information about chirality (could also be false, because this information is guessed anyway)


if kind_of_fragments == 'simple':
	#extract basic rings
	rules_standard = get_rules(password, username, 'simple')
else:
	rules_standard = get_rules(password, username, 'standard')

rules_after = get_rules(password, username, 'after')

neutralize_rules = chemistry.neutralize_get_rules()


print 'calculate frameworks'
for mol in smiles_file.GetOEMols():
	result_smiles = []
     	frameworks = []
     	OEFindRingAtomsAndBonds(mol)
     	OEAssignAromaticFlags(mol)
     	OEAssignFormalCharges(mol)
     	old_smi = chemistry.CanSmi(mol,iso,kek,chiral)
	old_smi = chemistry.neutralize_smi(old_smi,neutralize_rules,mol,debug)
	if debug:
        	print "==============================================="
		print old_smi, 'old_smi', mol.GetTitle()
	scaffolds = transform_framework(rules_standard,old_smi,mol.GetTitle(),debug, False)
	failed = False
	for i in range(0, len(scaffolds)):
                count_ring, ringlist = OEDetermineRingSystems(scaffolds[i])
		if debug:
                	print count_ring, ' count_ring'
 		list_of_scafs = [scaffolds[i]]
                if count_ring > 1:
			smiles_list = break_linkers(scaffolds[i],count_ring, ringlist)

			for smiles in smiles_list:

				final_smiles = transform_framework(rules_after,smiles,mol.GetTitle(),debug, False)
				final_smi = chemistry.CanSmi(final_smiles[0],iso,kek,chiral)
				if final_smi not in result_smiles:
					out_file.write(final_smi + '\t' + mol.GetTitle() + '\n')


                else:
			smi = chemistry.CanSmi(scaffolds[i],iso,kek,chiral)
			if smi not in result_smiles:
				result_smiles.append(smi)
				out_file.write(smi + '\t' + mol.GetTitle() + '\n')

conn.close ()
out_file_failed.close()
out_file.close()
print 'done'				
			

		
